# ==============================================================================
# Copyright (C) 2025 Evil0ctal
#
# This file is part of the Douyin_TikTok_Download_API project.
#
# This project is licensed under the Apache License 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at:
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import json
import os
import csv
from datetime import datetime
from typing import Dict, List, Any
import pandas as pd


class ExportManager:
    """数据导出管理器"""
    
    def __init__(self, export_dir: str = "./download/search_results"):
        """
        初始化导出管理器
        Args:
            export_dir: 导出文件目录
        """
        self.export_dir = export_dir
        self.ensure_export_dir()
    
    def ensure_export_dir(self):
        """确保导出目录存在"""
        if not os.path.exists(self.export_dir):
            os.makedirs(self.export_dir, exist_ok=True)
    
    def generate_filename(self, keyword: str, file_type: str = "json") -> str:
        """
        生成文件名
        Args:
            keyword: 搜索关键字
            file_type: 文件类型 (json/csv)
        Returns:
            生成的文件名
        """
        # 清理关键字，移除特殊字符
        clean_keyword = "".join(c for c in keyword if c.isalnum() or c in (' ', '-', '_')).strip()
        clean_keyword = clean_keyword.replace(' ', '_')
        
        # 生成时间戳
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        filename = f"{clean_keyword}_{timestamp}.{file_type}"
        return os.path.join(self.export_dir, filename)
    
    def export_to_json(self, data: Dict[str, Any], keyword: str) -> str:
        """
        导出数据到JSON文件
        Args:
            data: 要导出的数据
            keyword: 搜索关键字
        Returns:
            导出文件的路径
        """
        filepath = self.generate_filename(keyword, "json")
        
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        
        return filepath
    
    def export_to_csv(self, data: Dict[str, Any], keyword: str) -> str:
        """
        导出数据到CSV文件
        Args:
            data: 要导出的数据
            keyword: 搜索关键字
        Returns:
            导出文件的路径
        """
        filepath = self.generate_filename(keyword, "csv")
        
        # 准备CSV数据
        csv_data = []
        videos = data.get('videos', [])
        
        for video in videos:
            # 评论内容预览（取前3条评论，每条最多50字符）
            comments_preview = ""
            if video.get('comments'):
                preview_comments = []
                for comment in video['comments'][:3]:
                    text = comment.get('text', '')[:50]
                    if len(comment.get('text', '')) > 50:
                        text += "..."
                    preview_comments.append(text)
                comments_preview = " | ".join(preview_comments)
            
            csv_row = {
                "视频ID": video.get('aweme_id', ''),
                "视频描述": video.get('desc', ''),
                "作者昵称": video.get('author', {}).get('nickname', ''),
                "作者ID": video.get('author', {}).get('sec_user_id', ''),
                "点赞数": video.get('statistics', {}).get('digg_count', 0),
                "评论数": video.get('statistics', {}).get('comment_count', 0),
                "分享数": video.get('statistics', {}).get('share_count', 0),
                "播放数": video.get('statistics', {}).get('play_count', 0),
                "无水印视频链接": video.get('video_urls', {}).get('no_watermark', ''),
                "有水印视频链接": video.get('video_urls', {}).get('with_watermark', ''),
                "无水印高清视频": video.get('video_urls', {}).get('no_watermark_hq', ''),
                "有水印高清视频": video.get('video_urls', {}).get('with_watermark_hq', ''),
                "API链接": video.get('api_urls', {}).get('full', ''),
                "API链接精简": video.get('api_urls', {}).get('minimal', ''),
                "评论数量": len(video.get('comments', [])),
                "评论内容预览": comments_preview,
                "创建时间": video.get('create_time', 0)
            }
            csv_data.append(csv_row)
        
        # 使用pandas写入CSV，确保中文正确显示
        df = pd.DataFrame(csv_data)
        df.to_csv(filepath, index=False, encoding='utf-8-sig')
        
        return filepath
    
    def export_both_formats(self, data: Dict[str, Any], keyword: str) -> Dict[str, str]:
        """
        同时导出JSON和CSV格式
        Args:
            data: 要导出的数据
            keyword: 搜索关键字
        Returns:
            包含两个文件路径的字典
        """
        json_path = self.export_to_json(data, keyword)
        csv_path = self.export_to_csv(data, keyword)
        
        return {
            "json": json_path,
            "csv": csv_path
        }
    
    def get_export_summary(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        获取导出数据摘要
        Args:
            data: 数据
        Returns:
            摘要信息
        """
        videos = data.get('videos', [])
        total_videos = len(videos)
        total_comments = sum(len(video.get('comments', [])) for video in videos)
        total_likes = sum(video.get('statistics', {}).get('digg_count', 0) for video in videos)
        
        return {
            "total_videos": total_videos,
            "total_comments": total_comments,
            "total_likes": total_likes,
            "search_keyword": data.get('search_info', {}).get('keyword', ''),
            "search_time": data.get('search_info', {}).get('search_time', '')
        }
    
    def cleanup_old_files(self, days_old: int = 7):
        """
        清理过期文件
        Args:
            days_old: 删除多少天前的文件
        """
        if not os.path.exists(self.export_dir):
            return
        
        current_time = datetime.now()
        for filename in os.listdir(self.export_dir):
            filepath = os.path.join(self.export_dir, filename)
            if os.path.isfile(filepath):
                file_time = datetime.fromtimestamp(os.path.getctime(filepath))
                if (current_time - file_time).days > days_old:
                    try:
                        os.remove(filepath)
                        print(f"删除过期文件: {filename}")
                    except Exception as e:
                        print(f"删除文件失败 {filename}: {str(e)}")


# 创建全局导出管理器实例
export_manager = ExportManager()